import torch
import torch.nn as nn
from torch.optim import SGD

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel,self).__init__()
        self.fc1 = nn.Linear(10,10)
        self.activation = nn.ReLU()
        self.fc2 = nn.Linear(10,1)
        self.dropout = nn.Dropout(0.5)

    def forward(self,x):
        x = self.fc1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x


#训练模式
model = SimpleModel()
opt = SGD(model.parameters(),lr=0.01)
loss = nn.MSELoss()

train_data = torch.randn((1000,10))
train_labels = torch.randn((1000,1))
test_data = torch.randn((20,10))
test_labels = torch.randn((20,1))

model.train()
for epoch in range(50):
    _y = model(train_data)
    loss_value = loss(_y,train_labels)
    print('损失大小：',loss_value)
    loss_value.backward()
    opt.step()
    opt.zero_grad()

model.eval()

with torch.no_grad():
    test_outputs = model(test_data)
    test_loss = loss(test_outputs,test_labels)
    print('test loss is:',test_loss)